# -*- coding: utf-8 -*-
"""
@Time : 2023/4/1 6:17 PM
@Author: imamhusan
@Des: WebSocket管理器 一下有三种方式实现
"""
import json
from datetime import datetime
from typing import Dict
from fastapi import FastAPI, WebSocket
from sqlalchemy import func
from sqlalchemy.orm import Session, attributes, joinedload

from app.core.helpers import datetime_to_str
from app.database import engine
from app.models import Conversation
from app.schemas.conversations import MessageSM
from app.services.conversations import ConversationsServices
from app.services.upload import UploadFile

app = FastAPI()


# 单个客户端websocket实例
@app.websocket('/ws')
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        try:
            data = await websocket.receive_json()
            await websocket.send_json({"data": data, "msg": "发送json数据"})
        except:
            break


# 多个客户端将所有连接存储在字典中，更适用于多对多的连接场景
class WebSocketManager:
    def __init__(self):
        # 维护一个WebSocket连接的字典，key为client_id，value为WebSocket
        self.chat_page_users = []
        self.clients: Dict[int, WebSocket] = {}

    async def connect(self, websocket: WebSocket, client_id: int):
        # 接受WebSocket连接，将WebSocket连接添加到clients字典中
        await websocket.accept()
        self.clients[client_id] = websocket

    async def disconnect(self, client_id: int):
        # 断开WebSocket连接，并将连接从clients字典中删除
        del self.clients[client_id]

    async def chat_page_connect(self, client_id: int):
        # 进入聊天儿页面WebSocket连接
        self.chat_page_users.append(client_id)

    async def chat_page_disconnect(self, client_id: int):
        # 退出聊天儿页面断开WebSocket连接
        if client_id in self.chat_page_users:
            self.chat_page_users.remove(client_id)

    # 聊天儿功消息处理
    async def send_message(self, data: dict, sender_id: int, receiver_id: int):
        ws = self.clients.get(receiver_id)
        content = data.get("content")
        message_type = data.get("type")
        message_file_ext = data.get("ext")
        conversation = ConversationsServices.is_conversation(receiver_id, sender_id)
        if message_type == "text":
            # 如果是文本消息，则直接发送
            message_save = ConversationsServices.message_store(
                MessageSM(content=content, conversation_id=conversation.id, sender_id=sender_id,
                          receiver_id=receiver_id, message=content, type=message_type))
            if receiver_id is not None:
                # 判断 对方连接websocket
                if ws is not None:
                    # 判断 对方在聊天儿页面
                    if receiver_id in self.chat_page_users:
                        if message_save is not None:
                            data = message_save.obj_to_dict()
                            data = {"page": 'chat', "data": data, 'msg': '发送成功'}
                            await ws.send_json(data)
                    else:
                        if message_save:
                            # 未读消息
                            update = unread_message(sender_id, receiver_id)
                            unread_count = unread_message_count(receiver_id)
                            #
                            data = {"page": 'conversion_un_reared_message',
                                    "data": {"update": update.obj_to_dict()},
                                    'msg': '未读消息'}
                            tabBar = {"page": 'tabBar',
                                      "data": {"unread_count": unread_count},
                                      'msg': '未读消息总数'}
                            await ws.send_json(data)
                            await ws.send_json(tabBar)
                else:
                    if message_save:
                        # 未读消息
                        unread_message(sender_id, receiver_id)

            else:
                print('不能为空receiver_id')
        elif message_type in ["image", "audio"]:
            # 如果是图片消息，则读取文件并以二进制格式发送
            if receiver_id is not None:
                ws = self.clients.get(receiver_id)
                if ws is not None:
                    file = await UploadFile(file=content).createBase64File(ext=message_file_ext)
                    ConversationsServices.message_store(
                        MessageSM(content=content, conversation_id=conversation.id, sender_id=sender_id,
                                  receiver_id=receiver_id, message=file['file_path'], type=message_type))
                    return file
                else:
                    file = await UploadFile(file=content).createBase64File(ext=message_file_ext)
                    message_save = ConversationsServices.message_store(
                        MessageSM(content=content, conversation_id=conversation.id, sender_id=sender_id,
                                  receiver_id=receiver_id, message=file['file_path'], type=message_type))

                    if message_save:
                        # 未读消息
                        print(unread_message(sender_id, receiver_id))
        else:
            return print('数据格式错误')

    async def connect_message(self, data: dict, receiver_id: int):
        ws = self.clients.get(receiver_id)
        if ws is not None:
            try:
                await ws.send_json(data)
            except Exception as e:
                print(f"Failed to send message due to {e}")
        else:
            print('Websocket connection not found!')


# 未读消息
def unread_message(sender_id, receiver_id):
    conversation = ConversationsServices.is_conversation(sender_id, receiver_id)
    with Session(engine) as session:
        obj = session.query(Conversation).filter(Conversation.id == conversation.id)
        first = obj.first()
        obj.update({"unread_messages": 1 + first.unread_messages})
        session.commit()
    updated_obj = session.query(Conversation).filter(Conversation.id == conversation.id).options(
        joinedload(Conversation.participant),
        joinedload(Conversation.last_message)).first()
    return updated_obj


def unread_message_count(receiver_id):
    with Session(engine) as session:
        unread_total = session.query(func.sum(Conversation.unread_messages)).select_from(Conversation).filter_by(
            creator_id=receiver_id).scalar()
        return int(unread_total)
